import mitt from 'mitt';
import { type MaybeRef, reactive, ref, toValue, watch } from 'vue';
import { tryOnBeforeUnmount, tryOnScopeDispose, useDocumentVisibility } from '@peng_kai/kit/libs/vueuse';
import type { Simplify } from 'type-fest';
import { type TState, WSCloseCode, WebSocketClient } from './WebSocketClient';
import { Heartbeat } from './Heartbeat';

/**
  <div class="w-fll fixed left-0 top-0 z-10 bg-black p-1 text-white">
    <div>state: {{ appStore.socket.state }}</div>
    <div>delay: {{ appStore.socket.delay }}</div>
    <div>retryCount: {{ appStore.socket.retryCount }}</div>
    <div>retryCountdown: {{ appStore.socket.retryCountdown }}</div>
    <div>closeCode: {{ appStore.socket.closeCode }}</div>
    <div>closeMessage: {{ appStore.socket.closeMessage }}</div>
  </div>
 */

export type TUseWebSocketReturn = ReturnType<typeof useWebSocket>;

export function useWebSocket<P extends Record<string, any>, C = keyof P['Cmd']>(url: MaybeRef<string>, proto: P) {
  const emitter = mitt<Record<C & string, any>>();
  const heartbeat = new Heartbeat({
    ping: () => {
      const body = new proto.MessageBody();
      body.Cmd = proto.Cmd.PING;
      wsc.send(proto.MessageBody.encode(body).finish());
      return true;
    },
    pong: (data) => {
      const body = proto.MessageBody.decode(new Uint8Array(data));
      const cmd = proto.Cmd[body.Cmd];
      return cmd === 'PONG';
    },
  }, 4 * 1000, 3 * 1000);
  const wsc = new WebSocketClient(toValue(url), heartbeat);
  const state = ref<TState>('closed');
  const delay = ref(-1);
  const retryCount = ref(0);
  const retryCountdown = ref(0);
  const closeCode = ref(0);
  const closeMessage = ref('');

  heartbeat.onDelayChange = (value) => {
    delay.value = value;
  };
  wsc.onStateChange = (value) => {
    state.value = value;
  };
  wsc.onConnectCountdown = (count, countdown) => {
    retryCount.value = count;
    retryCountdown.value = countdown;
  };
  wsc.onClose = (event) => {
    closeCode.value = event.code;
    closeMessage.value = event.reason;
  };
  wsc.onMessage = (event) => {
    const res = proto.MessageBody.decode(new Uint8Array(event.data));
    emitter.emit(proto.Cmd[res.Cmd], res);
  };

  /**
   * 发送消息
   * @param cmd - 消息命令类型
   * @param encoder - 编码器类型
   * @param generator - 生成要发送的数据的函数
   */
  function send<QP extends keyof P>(
    cmd: GetReqMsgType<C>,
    encoder?: GetReqParser<QP>,
    generator?: <D extends InstanceType<P[QP]>>(data: D) => D,
  ) {
    const body = new proto.MessageBody();
    body.Cmd = proto.Cmd[cmd];

    if (encoder && generator) {
      const dataParser = new proto[encoder]();
      generator(dataParser);
      body.Data = proto[encoder].encode(dataParser).finish();
    }

    return wsc.send(proto.MessageBody.encode(body).finish());
  }

  /**
   * 接收服务器发送的消息，并通过指定的解码器对消息进行解码。
   * @param cmd 消息命令类型
   * @param decoder 解码器类型
   * @param listener 消息接收回调函数，接收解码后的数据和响应对象作为参数
   */
  function receive<PP extends keyof P>(
    cmd: GetRespMsgType<C>,
    decoder: GetRespParser<PP> | undefined,
    listener: (data: GetMessageBody<P, PP>) => void,
  ) {
    const _listener = (body: any) => {
      body.Data = body.Data instanceof Uint8Array ? proto[decoder ?? '']?.decode(body.Data) : body.Data;
      listener(body);
    };
    emitter.on(cmd, _listener);

    // 移除监听
    const off = () => {
      emitter.off(cmd, _listener);
    };

    tryOnBeforeUnmount(off);
    tryOnScopeDispose(off);

    return off;
  }

  function connect() {
    return wsc.connect();
  }
  function disconnect() {
    return wsc.close();
  }

  if (!import.meta.env.DEV) {
    receive(<any>'NEW_USER_CLIENT_RESP', undefined, () => {
      wsc.close(WSCloseCode.NEW_USER_CLIENT);
    });
  }

  window.addEventListener('beforeunload', () => {
    wsc.close();
  });

  watch(useDocumentVisibility(), (v) => {
    v === 'visible' && wsc.state === 'closed' && connect();
  });

  return reactive({ state, delay, retryCount, retryCountdown, closeCode, closeMessage, send, receive, connect, disconnect });
}

type GetReqMsgType<D> = D extends `${string}_REQ` ? D : never;
type GetReqParser<D> = Exclude<D, 'Cmd'> & string;
type GetRespMsgType<D> = D extends `${string}_RESP` ? D : never;
type GetRespParser<D> = Exclude<D, 'Cmd'> & string;

type GetMessageBody<P, K> = P extends { MessageBody: { decode: (...args: any) => infer B } }
  ? Simplify<
    Omit<B, 'Data'>
    & { Data: P extends { [I in K & string]: { decode: (...args: any) => infer D } } ? D : any }
  >
  : any;
